import os
import base64
from typing import Optional, List, Dict
from mcp.server.fastmcp import FastMCP
from zai import ZhipuAiClient


def get_api_key() -> str:
    """Get API key from environment variables"""
    api_key = os.getenv("ZHIPUAI_API_KEY")
    if not api_key:
        raise ValueError("ZHIPUAI_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


def _convert_local_image_to_base64(image_path: str) -> str:
    """
    将本地图片文件转换为base64格式

    参数:
        image_path: 本地图片文件路径

    返回:
        base64编码的图片数据

    异常:
        ValueError: 文件不存在或无法读取
    """
    if not os.path.exists(image_path):
        raise ValueError(f"本地图片文件不存在: {image_path}")
    
    try:
        with open(image_path, "rb") as image_file:
            image_data = image_file.read()
            base64_data = base64.b64encode(image_data).decode("utf-8")
            return base64_data
    except Exception as e:
        raise ValueError(f"无法读取本地图片文件 {image_path}: {str(e)}")


def format_multimodal_response(response_data: Dict) -> str:
    """
    格式化多模态响应数据

    参数:
        response_data: 多模态原始响应数据

    返回:
        格式化后的响应字符串
    """
    if "choices" not in response_data or not response_data["choices"]:
        return "API响应格式错误"

    message = response_data["choices"][0].get("message", {})
    content = message.get("content", "")

    return content if content else "响应内容为空"


mcp = FastMCP("ZhipuAI Multimodal MCP Server")


@mcp.tool()
async def understand_images(
    image_urls: List[str], question: str = "请描述这些图片的内容"
) -> str:
    """
    理解图片内容（使用智谱AI glm-4.5v模型）

    参数:
        image_urls: 图片URL列表或本地文件路径列表，例如：
                   ["https://example.com/image1.jpg", "/path/to/local/image.jpg"]
        question: 图片相关问题，默认值："请描述这些图片的内容"

    返回:
        图片理解结果字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        Exception: API调用失败
    """
    if not image_urls:
        raise ValueError("必须提供至少一个图片URL")

    api_key = get_api_key()
    client = ZhipuAiClient(api_key=api_key)

    # 构建内容数组
    content = []
    for url in image_urls:
        # 判断是HTTP URL还是本地文件路径
        if url.startswith(('http://', 'https://')):
            # HTTP URL，直接使用
            content.append({"type": "image_url", "image_url": {"url": url}})
        else:
            # 本地文件路径，转换为base64
            try:
                base64_data = _convert_local_image_to_base64(url)
                content.append({"type": "image_url", "image_url": {"url": base64_data}})
            except ValueError as e:
                raise ValueError(f"本地图片处理失败: {str(e)}")

    content.append({"type": "text", "text": question})

    try:
        response = client.chat.completions.create(
            model="glm-4.5v",
            messages=[{"role": "user", "content": content}],
            thinking={"type": "enabled"},
        )

        # 将SDK响应转换为字典格式
        response_dict = {
            "choices": [{"message": {"content": response.choices[0].message.content}}]
        }

        return format_multimodal_response(response_dict)
    except Exception as e:
        raise Exception(f"图片理解失败: {str(e)}")


@mcp.tool()
async def understand_video(
    video_url: str, question: str = "请描述这个视频的内容"
) -> str:
    """
    理解视频内容（使用智谱AI glm-4.5v模型）

    参数:
        video_url: 视频URL，例如：["https://example.com/video.mp4"]
        question: 视频相关问题，默认值："请描述这个视频的内容"

    返回:
        视频理解结果字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        Exception: API调用失败
    """
    if not video_url:
        raise ValueError("必须提供视频URL")

    api_key = get_api_key()
    client = ZhipuAiClient(api_key=api_key)

    try:
        response = client.chat.completions.create(
            model="glm-4.5v",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "video_url", "video_url": {"url": video_url}},
                        {"type": "text", "text": question},
                    ],
                }
            ],
            thinking={"type": "enabled"},
        )

        # 将SDK响应转换为字典格式
        response_dict = {
            "choices": [{"message": {"content": response.choices[0].message.content}}]
        }

        return format_multimodal_response(response_dict)
    except Exception as e:
        raise Exception(f"视频理解失败: {str(e)}")


@mcp.tool()
async def understand_files(
    file_urls: List[str], question: str = "请分析这些文件的内容"
) -> str:
    """
    理解文件内容（使用智谱AI glm-4.5v模型）

    参数:
        file_urls: 文件URL列表，例如：["https://example.com/file1.pdf", "https://example.com/file2.txt"]
        question: 文件相关问题，默认值："请分析这些文件的内容"

    返回:
        文件理解结果字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        Exception: API调用失败
    """
    if not file_urls:
        raise ValueError("必须提供至少一个文件URL")

    api_key = get_api_key()
    client = ZhipuAiClient(api_key=api_key)

    # 构建内容数组
    content = []
    for url in file_urls:
        content.append({"type": "file_url", "file_url": {"url": url}})

    content.append({"type": "text", "text": question})

    try:
        response = client.chat.completions.create(
            model="glm-4.5v",
            messages=[{"role": "user", "content": content}],
            thinking={"type": "enabled"},
        )

        # 将SDK响应转换为字典格式
        response_dict = {
            "choices": [{"message": {"content": response.choices[0].message.content}}]
        }

        return format_multimodal_response(response_dict)
    except Exception as e:
        raise Exception(f"文件理解失败: {str(e)}")


@mcp.tool()
async def multimodal_analysis(
    image_urls: Optional[List[str]] = None,
    video_url: Optional[str] = None,
    file_urls: Optional[List[str]] = None,
    question: str = "请分析提供的多媒体内容",
) -> str:
    """
    综合多模态分析（使用智谱AI glm-4.5v模型）

    参数:
        image_urls: 可选的图片URL列表
        video_url: 可选的视频URL
        file_urls: 可选的文件URL列表
        question: 分析问题，默认值："请分析提供的多媒体内容"

    返回:
        多模态分析结果字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        Exception: API调用失败
    """
    if not image_urls and not video_url and not file_urls:
        raise ValueError("必须提供至少一种多媒体内容（图片、视频或文件）")

    api_key = get_api_key()
    client = ZhipuAiClient(api_key=api_key)

    # 构建内容数组
    content = []

    if image_urls:
        for url in image_urls:
            content.append({"type": "image_url", "image_url": {"url": url}})

    if video_url:
        content.append({"type": "video_url", "video_url": {"url": video_url}})

    if file_urls:
        for url in file_urls:
            content.append({"type": "file_url", "file_url": {"url": url}})

    content.append({"type": "text", "text": question})

    try:
        response = client.chat.completions.create(
            model="glm-4.5v",
            messages=[{"role": "user", "content": content}],
            thinking={"type": "enabled"},
        )

        # 将SDK响应转换为字典格式
        response_dict = {
            "choices": [{"message": {"content": response.choices[0].message.content}}]
        }

        return format_multimodal_response(response_dict)
    except Exception as e:
        raise Exception(f"多模态分析失败: {str(e)}")


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
